/**
 * 
 */
package org.apache.lucene.search;

import java.io.IOException;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.grouping.SearchGroup;
import org.apache.lucene.search.grouping.TermSecondPassGroupingCollector;

/**
 * @author Liuye
 * 
 */
public class B2BTermSecondPassGroupingCollector2 extends
		TermSecondPassGroupingCollector {
	private Map<String, Integer> groupingDistinctHitCount;
	private FieldCache.StringIndex index;
	private String groupField;
	private int maxDocsPerGroup;

	public B2BTermSecondPassGroupingCollector2(String groupField,
			Collection<SearchGroup<String>> groups, Sort groupSort,
			Sort withinGroupSort, int maxDocsPerGroup, boolean getScores,
			boolean getMaxScores, boolean fillSortFields) throws IOException {
		super(groupField, groups, groupSort, withinGroupSort, maxDocsPerGroup,
				getScores, getMaxScores, fillSortFields);
		this.groupField = groupField;
		this.maxDocsPerGroup = maxDocsPerGroup;
		this.groupingDistinctHitCount = new HashMap<String, Integer>();
	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see
	 * org.apache.lucene.search.grouping.AbstractSecondPassGroupingCollector
	 * #collect(int)
	 */
	@Override
	public void collect(int doc) throws IOException {
		super.collect(doc);
		String groupValue = getDocGroupValue(doc);
		Integer groupCount = groupingDistinctHitCount.get(groupValue);
		if (groupCount == null || groupCount < maxDocsPerGroup) {
			groupCount = groupCount == null ? 1 : groupCount++;
			groupingDistinctHitCount.put(groupValue, groupCount);
		}
	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see org.apache.lucene.search.grouping.TermSecondPassGroupingCollector#
	 * setNextReader(org.apache.lucene.index.IndexReader, int)
	 */
	@Override
	public void setNextReader(IndexReader reader, int docBase)
			throws IOException {
		super.setNextReader(reader, docBase);
		index = FieldCache.DEFAULT.getStringIndex(reader, groupField);
	}

	protected String getDocGroupValue(int doc) {
		final int ord = index.order[doc];
		return ord == 0 ? null : index.lookup[ord];
	}

	public int getGroupingDistinctHitCount() {
		int count = 0;
		for (Map.Entry<String, Integer> entry : groupingDistinctHitCount.entrySet()) {
			count += entry.getValue();
		}
		return count;
	}

}
